package com.ajocer.springbootinit.utils;

import io.grpc.ManagedChannel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.util.Arrays;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * gRPC连接监控工具
 * 监控gRPC连接健康状态和检测潜在的连接泄漏
 */
@Component
@Slf4j
public class GrpcConnectionMonitor {

    @Value("${grpc.monitor.enabled:true}")
    private boolean monitorEnabled;

    private final AtomicInteger grpcThreadCount = new AtomicInteger(0);
    private final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();

    /**
     * 定期检查gRPC线程数量
     * 如果发现gRPC线程数量异常增长，可能存在连接泄漏
     */
    @Scheduled(fixedRate = 60000) // 每分钟执行一次
    public void monitorGrpcThreads() {
        if (!monitorEnabled) {
            return;
        }

        ThreadInfo[] threadInfos = threadMXBean.dumpAllThreads(false, false);
        long grpcThreads = 0;
        for (ThreadInfo threadInfo : threadInfos) {
            if (threadInfo.getThreadName().contains("grpc-")) {
                grpcThreads++;
            }
        }

        int previousCount = grpcThreadCount.getAndSet((int) grpcThreads);

        // 如果gRPC线程数量增长超过10个，记录警告日志
        if (grpcThreads > previousCount + 10) {
            log.warn("检测到gRPC线程数量异常增长: 从 {} 增加到 {}, 可能存在连接泄漏", 
                    previousCount, grpcThreads);
        }

        // 记录当前gRPC线程数量
        log.debug("当前gRPC线程数量: {}", grpcThreads);
    }

    /**
     * 安全关闭gRPC通道
     * @param channel 要关闭的gRPC通道
     */
    public static void shutdownChannel(ManagedChannel channel) {
        if (channel != null && !channel.isShutdown()) {
            try {
                // 尝试优雅关闭，等待5秒
                channel.shutdown();
                if (!channel.awaitTermination(5, TimeUnit.SECONDS)) {
                    // 如果5秒后仍未关闭，强制关闭
                    channel.shutdownNow();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                // 如果被中断，强制关闭
                channel.shutdownNow();
            }
        }
    }
} 